Skip to content

Prefix sum (scan) implementation (C2V sync).#101

Merged
zouzias merged 4 commits into
huawei-csl:mainfrom
vloncar:scan
Apr 16, 2026
Merged

Prefix sum (scan) implementation (C2V sync).#101
zouzias merged 4 commits into
huawei-csl:mainfrom
vloncar:scan

Conversation

@vloncar

@vloncar vloncar commented Apr 14, 2026

Copy link
Copy Markdown
Collaborator

An implementation of prefix sum (scan) algorithm, based on https://arxiv.org/abs/2505.15112v1. Currently, only single core algorithm is implemented (ScanU from the paper).

The implementation requires CUBE <-> AIV synchronization, for which I made a custom sync implementation inspired by the one in #65. I had to make changes to sync on PIPE_FIX and MTE3, as the TSync_Custom.hpp implementation suggests. Perhaps the issue observed in #65 has since been resolved.

@zouzias

zouzias commented Apr 16, 2026

Copy link
Copy Markdown
Collaborator

Tests pass on 910B4.

python run_scan.py
/home/zouzias/github-repos/pto-kernels/venv/lib/python3.10/site-packages/torch_npu/utils/collect_env.py:58: UserWarning: Warning: The /usr/local/Ascend/cann-8.5.0 owner does not match the current owner.
  warnings.warn(f"Warning: The {path} owner does not match the current owner.")
/home/zouzias/github-repos/pto-kernels/venv/lib/python3.10/site-packages/torch_npu/utils/collect_env.py:58: UserWarning: Warning: The /usr/local/Ascend/cann-8.5.0/x86_64-linux/ascend_ops_install.info owner does not match the current owner.
  warnings.warn(f"Warning: The {path} owner does not match the current owner.")
/home/zouzias/github-repos/pto-kernels/venv/lib/python3.10/site-packages/torch_npu/utils/collect_env.py:58: UserWarning: Warning: The /usr/local/Ascend/cann-8.5.0 owner does not match the current owner.
  warnings.warn(f"Warning: The {path} owner does not match the current owner.")
/home/zouzias/github-repos/pto-kernels/venv/lib/python3.10/site-packages/torch_npu/utils/collect_env.py:58: UserWarning: Warning: The /usr/local/Ascend/cann-8.5.0/x86_64-linux/ascend_ops_install.info owner does not match the current owner.
  warnings.warn(f"Warning: The {path} owner does not match the current owner.")
Testing NPU scan kernel: tile_size=16x16, total_len=256 (1 tiles)
All results matched. Scan test passed successfully.

Testing NPU scan kernel: tile_size=16x16, total_len=4096 (16 tiles)
All results matched. Scan test passed successfully.

Testing NPU scan kernel: tile_size=16x16, total_len=16384 (64 tiles)
All results matched. Scan test passed successfully.

Testing NPU scan kernel: tile_size=16x16, total_len=25600 (100 tiles)
All results matched. Scan test passed successfully.

Testing NPU scan kernel: tile_size=32x32, total_len=1024 (1 tiles)
All results matched. Scan test passed successfully.

Testing NPU scan kernel: tile_size=32x32, total_len=16384 (16 tiles)
All results matched. Scan test passed successfully.

Testing NPU scan kernel: tile_size=32x32, total_len=65536 (64 tiles)
All results matched. Scan test passed successfully.

Testing NPU scan kernel: tile_size=32x32, total_len=102400 (100 tiles)
All results matched. Scan test passed successfully.

Testing NPU scan kernel: tile_size=64x64, total_len=4096 (1 tiles)
All results matched. Scan test passed successfully.

Testing NPU scan kernel: tile_size=64x64, total_len=65536 (16 tiles)
All results matched. Scan test passed successfully.

Testing NPU scan kernel: tile_size=64x64, total_len=262144 (64 tiles)
All results matched. Scan test passed successfully.

Testing NPU scan kernel: tile_size=64x64, total_len=409600 (100 tiles)
All results matched. Scan test passed successfully.

Testing NPU scan kernel: tile_size=128x128, total_len=16384 (1 tiles)
All results matched. Scan test passed successfully.

Testing NPU scan kernel: tile_size=128x128, total_len=262144 (16 tiles)
All results matched. Scan test passed successfully.

Testing NPU scan kernel: tile_size=128x128, total_len=1048576 (64 tiles)
All results matched. Scan test passed successfully.

Testing NPU scan kernel: tile_size=128x128, total_len=1638400 (100 tiles)
All results matched. Scan test passed successfully.

@zouzias zouzias left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@vloncar , please add a README.md file like below and merge. Nice MR!

# Single core prefix sum (scan)

An implementation of prefix sum (scan) algorithm, based on https://arxiv.org/abs/2505.15112v1. Currently, only single core algorithm is implemented (ScanU from the paper).

Usage:

```bash

# Optional
export PTO_LIB_PATH=${ASCEND_TOOLKIT_HOME}  # reuse CANN 8.5.0 headers

# Run scan tests
python ./run_scan.py

@zouzias zouzias marked this pull request as ready for review April 16, 2026 03:57
@zouzias zouzias requested a review from gioelegott April 16, 2026 04:01
Comment thread examples/jit_cpp/scan/run_scan.py Outdated
@zouzias zouzias merged commit bf6c9bc into huawei-csl:main Apr 16, 2026
15 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants